// The service is a bootstrap for a local system process to start
// when the computer starts. The service shuts itself down in 30 seconds.
//
// The service can be started in one of two modes:
//   * As a regular service, typically at system startup.
//   * As a COM service, typically by an PCAgent client using ILongkeyPCAgentCore.
// The COM case is distinguished from the regular service case by registering a
// ServiceParameters command line in the AppID registration.
//
// In all cases, the service initializes COM, and allows for ILongkeyPCAgentCore
// clients to connect to the service. In the regular service case, the service
// shuts down after a small idle check timeout, provided that there are no COM
// clients connected. In the COM server case, and in the case where there are
// COM clients connected in the regular service case, the service will shut down
// when the last client disconnects.
//
// To be exact, the service will initiate shutdown in all cases when the ATL
// module count drops to zero.
//
// ATL does not allow for directly reusing the delayed COM shutdown mechanism
// available for Local servers. The assumption likely being that services run
// forever. Since we do not want our service to run forever, we override some
// of the functions to get the same effect.

#ifndef LONGKEY_SERVICE_SERVICE_MAIN_H_
#define LONGKEY_SERVICE_SERVICE_MAIN_H_

#include <atlbase.h>
#include <atlcom.h>
#include "../base/basictypes.h"
#include "../base/atlregmapex.h"
#include "../base/constants.h"
#include "../base/reg_key.h"
#include "../common/command_line.h"
#include "../common/const_cmd_line.h"
#include "../pcagentres/resource.h"
#include "../pcagent/non_localized_resource.h"

#pragma warning(push)
// C4640: construction of local static object is not thread-safe
#pragma warning(disable : 4640)

class PCAgentServiceMode
{
public:
	static CommandLineMode commandline_mode();
	static CString reg_name();
	static CString default_name();
	static DWORD service_start_type();
	static _ATL_OBJMAP_ENTRY* object_map();
	static bool allow_access_from_medium();
	static CString app_id_string();
	static CString GetCurrentServiceName();
	static HRESULT PreMessageLoop();
};

const TCHAR* const kServiceFileName              = kPCAgentShellFileName;

template <typename T>
class ServiceModule : public CAtlServiceModuleT<ServiceModule<T>, IDS_SERVICE_NAME>
{
public:
	typedef CAtlServiceModuleT<ServiceModule, IDS_SERVICE_NAME> Base;

	DECLARE_REGISTRY_APPID_RESOURCEID_EX(IDR_LONGKEY_PCAGENT_SERVICE_APPID, T::app_id_string());

	BEGIN_REGISTRY_MAP()
		REGMAP_ENTRY(_T("DESCRIPTION"),   _T("ServiceModule"))
		REGMAP_ENTRY(_T("FILENAME"),      kServiceFileName)
	END_REGISTRY_MAP()

	ServiceModule() : service_thread_(NULL), is_service_com_server_(false)
	{
		_tcscpy_s(m_szServiceName, 255, T::GetCurrentServiceName());
	}

	~ServiceModule()
	{
		ASSERT1(!service_thread_);

		// ServiceModule is typically created on the stack. We cannot reset the
		// _pAtlModule here, because objects are destroyed at destructor unwind
		// time, and require access to the module.
	}

	HRESULT InitializeSecurity() throw()
	{
		return InitializeServerSecurity(T::allow_access_from_medium());
	}

	void ServiceMain(DWORD argc, LPTSTR* argv) throw()
	{
		ASSERT1(argc <= 2);
		is_service_com_server_ = argc == 2 && !CString(argv[1]).CompareNoCase(kCmdLineServiceComServer);

		Base::ServiceMain(argc, argv);
	}

	HRESULT RegisterClassObjects(DWORD class_context, DWORD flags) throw()
	{
		for (_ATL_OBJMAP_ENTRY* objmap_entry = T::object_map();
			objmap_entry->pclsid != NULL;
			objmap_entry++)
		{
			HRESULT hr = objmap_entry->RegisterClassObject(class_context, flags);
			if (FAILED(hr))
				return hr;
		}
		return S_OK;
	}

	HRESULT RevokeClassObjects() throw()
	{
		for (_ATL_OBJMAP_ENTRY* objmap_entry = T::object_map();
			objmap_entry->pclsid != NULL;
			objmap_entry++)
		{
			HRESULT hr = objmap_entry->RevokeClassObject();
			if (FAILED(hr))
				return hr;
		}
		return S_OK;
	}

	HRESULT RegisterServer(BOOL, const CLSID* = NULL) throw()
	{
		for (_ATL_OBJMAP_ENTRY* objmap_entry = T::object_map();
			objmap_entry->pclsid != NULL;
			objmap_entry++)
		{
			HRESULT hr = objmap_entry->pfnUpdateRegistry(TRUE);
			if (FAILED(hr))
				return hr;

			hr = AtlRegisterClassCategoriesHelper(*objmap_entry->pclsid,
				objmap_entry->pfnGetCategoryMap(), TRUE);
			if (FAILED(hr))
				return hr;
		}
		return S_OK;
	}

	HRESULT UnregisterServer(BOOL, const CLSID* = NULL) throw()
	{
		for (_ATL_OBJMAP_ENTRY* objmap_entry = T::object_map();
			objmap_entry->pclsid != NULL;
			objmap_entry++) {
				HRESULT hr = AtlRegisterClassCategoriesHelper(*objmap_entry->pclsid,
					objmap_entry->pfnGetCategoryMap(),
					FALSE);
				if (FAILED(hr))
					return hr;

				hr = objmap_entry->pfnUpdateRegistry(FALSE);
				if (FAILED(hr))
					return hr;
		}

		return S_OK;
	}


	HRESULT RegisterCOMService()
	{
		HRESULT hr = UpdateRegistryAppId(TRUE);
		if (FAILED(hr))
			return hr;

		RegKey key_app_id;
		hr = key_app_id.Open(HKEY_CLASSES_ROOT, _T("AppId"), KEY_WRITE);
		if (FAILED(hr))
			return hr;

		RegKey key;
		hr = key.Create(key_app_id.Key(), GetAppIdT());
		if (FAILED(hr))
			return hr;

		// m_szServiceName is set in the constructor
		hr = key.SetValue(_T("LocalService"), m_szServiceName);
		if (FAILED(hr))
			return hr;

		// The SCM will pass this switch to ServiceMain() during COM activation.
		hr = key.SetValue(_T("ServiceParamters"), kCmdLineServiceComServer);
		if (FAILED(hr))
			return hr;

		return RegisterServer(FALSE);
	}

	HRESULT UnregisterCOMService()
	{
		HRESULT hr = UnregisterServer(FALSE);
		if (FAILED(hr))
			return hr;

		return UpdateRegistryAppId(FALSE);
	}

	HRESULT PreMessageLoop(int show_cmd)
	{
		UNREFERENCED_PARAMETER(show_cmd);

		m_dwThreadID = ::GetCurrentThreadId();
		service_thread_ = ::OpenThread(SYNCHRONIZE, FALSE, m_dwThreadID);

		if (is_service_com_server_)
			return InitializeCOMServer();

		// This is the regular service case. Call T::PreMessageLoop() and exit;
		SetServiceStatus(SERVICE_RUNNING);

		HRESULT hr = T::PreMessageLoop();
		if (FAILED(hr))
		{

		}

		SetServiceStatus(SERVICE_STOP_PENDING);

		// S_FALSE is returned to exit the service immediately.
		return S_FALSE;
	}

	HRESULT PostMessageLoop()
	{
		if (!is_service_com_server_)
			return S_OK;
		return Base::PostMessageLoop();
	}

	int Main(int show_cmd)
	{
		if (CAtlBaseModule::m_bInitFailed)
			return -1;

		return static_cast<int>(Start(show_cmd));
	}

	// This is cloned from CAtlExeModuleT.Lock(). The one difference is the call
	// to ::CoAddRefServerProcess(). See the description for Unlock() below for
	// further information.
	virtual LONG Lock() throw()
	{
		::CoAddRefServerProcess();
		LONG retval = CComGlobalsThreadModel::Increment(&m_nLockCnt);
		return retval;
	}

	// This is cloned from CAtlExeModuleT.Unlock(). The differences are:
	//
	// * the call to ::CoReleaseServerProcess(), to ensure that the class
	// factories are suspended once the lock count drops to zero. This fixes a
	// a race condition where an activation request could come in in the middle
	// of shutting down. This shutdown mechanism works with free threaded servers.
	//
	// There are race issues with the ATL  delayed shutdown mechanism, hence the
	// associated code has been eliminated, and we have an assert to make sure
	// m_bDelayShutdown is not set.
	//
	// * the call to "OnStop()" instead of the "::PostThreadMessage(m_dwMainThre".
	// OnStop() correctly sets the service status to SERVICE_STOP_PENDING, and
	// posts a WM_QUIT to the service thread.
	virtual LONG Unlock() throw()
	{
		ASSERT1(!m_bDelayShutdown);

		::CoReleaseServerProcess();
		LONG retval = CComGlobalsThreadModel::Decrement(&m_nLockCnt);

		if (retval == 0)
			OnStop();

		return retval;
	}

	// This is cloned from CAtlExeModuleT.MonitorShutdown(). The only difference
	// is the call to "OnStop()" instead of the
	// "::PostThreadMessage(m_dwMainThreadID".
	void MonitorShutdown() throw()
	{
		while (true) {
			::WaitForSingleObject(m_hEventShutdown, INFINITE);

			DWORD wait = 0;
			do
			{
				m_bActivity = false;
				wait = ::WaitForSingleObject(m_hEventShutdown, m_dwTimeOut);
			} while (wait == WAIT_OBJECT_0);

			if (!m_bActivity && m_nLockCnt == 0)
			{
				::CoSuspendClassObjects();
				if (m_nLockCnt == 0)
					break;
			}
		}

		::CloseHandle(m_hEventShutdown);
		OnStop();
	}

	// This is cloned from CAtlExeModuleT.StartMonitor().
	HANDLE StartMonitor() throw()
	{
		m_hEventShutdown = ::CreateEvent(NULL, false, false, NULL);
		if (m_hEventShutdown == NULL)
			return NULL;

		DWORD thread_id(0);
		HANDLE monitor = ::CreateThread(NULL, 0, MonitorProc, this, 0, &thread_id);
		if (monitor == NULL)
			::CloseHandle(m_hEventShutdown);

		return monitor;
	}

	//// This is cloned from CAtlExeModuleT.MonitorProc().
	//static DWORD WINAPI MonitorProc(void* pv) throw()
	//{
	//	ServiceModule* service_module = static<ServiceModule*>(pv);

	//	service_module->MonitorShutdown();
	//	return 0;
	//}


private:
	// Should only be called from the client, such as during setup, since it only
	// supports one language.
	static CString GetServiceDescription()
	{
		CString company_name;
		VERIFY1(company_name.LoadString(IDS_FRIENDLY_COMPANY_NAME));
		CString description;
		description.FormatMessage(IDS_SERVICE_DESCRIPTION, company_name);
		return description;
	}

	HRESULT InitializeCOMServer()
	{
		// Initialize COM security right at the beginning, because Worker
		// initialization can cause interface marshaling. The first few lines below
		// are adapted from the beginning of CAtlServiceModuleT::PreMessageLoop().
		ASSERT1(m_bService);
		HRESULT hr = InitializeSecurity();
		if (FAILED(hr))
			return hr;

		DisableCOMExceptionHandling();


		return CAtlExeModuleT<ServiceModule>::PreMessageLoop(SW_HIDE);
	}

	// When Start executes, it blocks on StartServiceCtrlDispatcher.
	// Internally, the SCM creates a service thread which starts executing code
	// specified by SERVICE_TABLE_ENTRY. The ATL code then calls PreMessageLoop
	// and PostMessageLoop on this thread. When the service stops, the execution
	// flow returns from StartServiceCtrlDispatcher and the main thread blocks
	// waiting on the service thread to exit.
	// Before synchronizing the main thread and the service thread, race condition
	// resulted in http://b/1134747.
	HRESULT Start(int show_cmd)
	{
		UNREFERENCED_PARAMETER(show_cmd);

		SERVICE_TABLE_ENTRY st[] =
		{
			{ m_szServiceName, Base::_ServiceMain },
			{ NULL, NULL }
		};

		m_status.dwWin32ExitCode = 0;
		if (!::StartServiceCtrlDispatcher(st))
			m_status.dwWin32ExitCode = ::GetLastError();

		if (service_thread_)
		{
			DWORD result(::WaitForSingleObject(service_thread_, INFINITE));
			ASSERT1(result == WAIT_OBJECT_0);
			::CloseHandle(service_thread_);
			service_thread_ = NULL;
		}

		return m_status.dwWin32ExitCode;
	}

	HANDLE service_thread_;			// The service thread provided by the SCM
	bool is_service_com_server_;	// True if the service is being invoked by COM.

	// Service fail over constants.
	//
	// Time after which the SCM reset the failure count to zero if there are
	// no failures.
	static const DWORD kResetPeriodSec		= 60 * 60 * 24;	// 1 day.

	DISALLOW_EVIL_CONSTRUCTORS(ServiceModule);
};

#pragma warning(pop)


typedef ServiceModule<PCAgentServiceMode> PCAgentServiceModule;

#endif